package org.example.semantic;

import com.google.common.collect.Sets;
import org.antlr.v4.runtime.tree.ParseTree;
import org.antlr.v4.runtime.tree.TerminalNode;

import java.util.*;
import java.util.regex.Pattern;

import static org.example.semantic.SemanticTokensProvider.TokenModifier.DECLARATION;
import static org.example.semantic.SemanticTokensProvider.TokenModifier.NONE;
import static org.example.semantic.SemanticTokensProvider.TokenType.*;


public class TokenIdentifyVisitor extends com.neuronbit.parser.MySqlParserBaseVisitor {
    private final SemanticTokensProvider semanticTokensProvider;
    private final Map<TerminalNode, SemanticInfo> identifiedTokens = new HashMap<>();
    private final Set<String> operators  = Sets.newHashSet(",", "(", ")", "=", ".", ";", "+", "-", "*", "\\");

    public TokenIdentifyVisitor(SemanticTokensProvider semanticTokensProvider) {
        this.semanticTokensProvider = semanticTokensProvider;
    }

    @Override
    public Object visitTableName(com.neuronbit.parser.MySqlParser.TableNameContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(PROPERTY, DECLARATION));
        }
        return super.visitTableName(ctx);
    }

    @Override
    public Object visitColumnName(com.neuronbit.parser.MySqlParser.ColumnNameContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(PROPERTY, DECLARATION));
        }
        return super.visitColumnName(ctx);
    }

    @Override
    public Object visitPartitionName(com.neuronbit.parser.MySqlParser.PartitionNameContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(VARIABLE, NONE));
        }
        return super.visitPartitionName(ctx);
    }

    @Override
    public Object visitDimensionDataType(com.neuronbit.parser.MySqlParser.DimensionDataTypeContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(FUNCTION, NONE));
        }
        return super.visitDimensionDataType(ctx);
    }

    @Override
    public Object visitStringDataType(com.neuronbit.parser.MySqlParser.StringDataTypeContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(FUNCTION, NONE));
        }
        return super.visitStringDataType(ctx);
    }

    @Override
    public Object visitLongVarcharDataType(com.neuronbit.parser.MySqlParser.LongVarcharDataTypeContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(FUNCTION, NONE));
        }
        return super.visitLongVarcharDataType(ctx);
    }

    @Override
    public Object visitLongVarbinaryDataType(com.neuronbit.parser.MySqlParser.LongVarbinaryDataTypeContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(FUNCTION, NONE));
        }
        return super.visitLongVarbinaryDataType(ctx);
    }

    @Override
    public Object visitSimpleIndexDeclaration(com.neuronbit.parser.MySqlParser.SimpleIndexDeclarationContext ctx) {
        for (ParseTree child : ctx.children) {
            if (child instanceof com.neuronbit.parser.MySqlParser.UidContext) {
                Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(child);
                if (optional.isPresent()) {
                    TerminalNode node = optional.get();
                    identifiedTokens.put(node, new SemanticInfo(PROPERTY, NONE));
                }
                break;
            }
        }
        return super.visitSimpleIndexDeclaration(ctx);
    }

    @Override
    public Object visitStringLiteral(com.neuronbit.parser.MySqlParser.StringLiteralContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(STRING, NONE));
        }
        return super.visitStringLiteral(ctx);
    }

    @Override
    public Object visitIndexColumnName(com.neuronbit.parser.MySqlParser.IndexColumnNameContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(PROPERTY, NONE));
        }
        return super.visitIndexColumnName(ctx);
    }

    @Override
    public Object visitCharsetNameBase(com.neuronbit.parser.MySqlParser.CharsetNameBaseContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(ENUMMEMBER, NONE));
        }
        return super.visitCharsetNameBase(ctx);
    }

    @Override
    public Object visitEngineName(com.neuronbit.parser.MySqlParser.EngineNameContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(ENUMMEMBER, NONE));
        }
        return super.visitEngineName(ctx);
    }

    @Override
    public Object visitCollationName(com.neuronbit.parser.MySqlParser.CollationNameContext ctx) {
        Optional<TerminalNode> optional = getTerminalByTreeOrFirstChildDfs(ctx.getChild(0));
        if (optional.isPresent()) {
            TerminalNode node = optional.get();
            identifiedTokens.put(node, new SemanticInfo(ENUMMEMBER, NONE));
        }
        return super.visitCollationName(ctx);
    }

    @Override
    public Object visitTerminal(TerminalNode node) {
        if ("<EOF>".equals(node.getText())) {
            return super.visitTerminal(node);
        }
        if (identifiedTokens.containsKey(node)) {
            semanticTokensProvider.push(node, identifiedTokens.get(node).tokenType, identifiedTokens.get(node).tokenModifiers);
        } else if (operators.contains(node.getText())) {
            semanticTokensProvider.push(node, SemanticTokensProvider.TokenType.OPERATOR, SemanticTokensProvider.TokenModifier.DECLARATION);
        } else if (isComment(node)) {
            semanticTokensProvider.push(node, SemanticTokensProvider.TokenType.COMMENT, SemanticTokensProvider.TokenModifier.DECLARATION);
        } else if (isNumber(node)) {
            semanticTokensProvider.push(node, SemanticTokensProvider.TokenType.NUMBER, SemanticTokensProvider.TokenModifier.NONE);
        } else {
            semanticTokensProvider.push(node, SemanticTokensProvider.TokenType.KEYWORD, SemanticTokensProvider.TokenModifier.DECLARATION);
        }
        return super.visitTerminal(node);
    }

    static Pattern pattern = Pattern.compile("^[0-9]+$");
    public boolean isNumber(TerminalNode node) {
        return pattern.matcher(node.getText()).find();
    }

    public boolean isComment(TerminalNode node) {
        return (node.getText().startsWith("'") && node.getText().endsWith("'")) ||
                (node.getText().startsWith("\"") && node.getText().endsWith("\""));
    }

    public Optional<TerminalNode> getTerminalByTreeOrFirstChildDfs(ParseTree tree) {
        if (Objects.isNull(tree)) {
            return Optional.empty();
        }
        while (!(tree instanceof TerminalNode)) {
            if (tree.getChildCount() < 1) {
                return Optional.empty();
            }
            tree = tree.getChild(0);
        }
        return Optional.of((TerminalNode) tree);
    }
    
    private static class SemanticInfo {
        private int tokenType;
        private int tokenModifiers = 0;

        public SemanticInfo(int tokenType, int tokenModifiers) {
            this.tokenType = tokenType;
            this.tokenModifiers = tokenModifiers;
        }
    }
}
